from transformers import GPT2Tokenizer, GPT2LMHeadModel
from sentence_transformers import SentenceTransformer
import faiss
import os

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["TRANSFORMERS_VERBOSITY"] = "info"

# 初始化嵌入模型和生成模型
embedding_model_name = "all-MiniLM-L6-v2"
embedding_model = SentenceTransformer(embedding_model_name)

generation_model_name = "gpt2"
tokenizer_generation = GPT2Tokenizer.from_pretrained(generation_model_name)
generation_model = GPT2LMHeadModel.from_pretrained(generation_model_name)

# 模拟文档数据集
documents = [
    "人工智能是模仿人类智能的技术，可用于复杂任务。",
    "机器学习是一种通过数据自动改进性能的技术。",
    "深度学习是机器学习的子集，主要通过神经网络实现。",
    "自然语言处理使得计算机能够理解和生成文本。",
    "数据科学是一门从数据中提取知识的学科。",
    # 可以扩展文档数量
]

# 生成文档嵌入
doc_embeddings = embedding_model.encode(documents)

# 使用FAISS创建索引
dimension = doc_embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(doc_embeddings)


# 检索模块：基于查询文本返回最相关的上下文


def retrieve_context(query_text, top_k=2):
    query_embedding = embedding_model.encode([query_text])
    distances, indices = index.search(query_embedding, top_k)
    results = [documents[i] for i in indices[0]]
    return results


# 生成模块：基于检索上下文生成回答
def generate_answer(contexts, query):
    context_text = " ".join(contexts)  # 将多个上下文合并
    input_text = f"背景信息: {context_text}\n问题: {query}\n回答:"
    inputs = tokenizer_generation(input_text, return_tensors="pt")

    # 使用GPT-2生成回答
    output_sequences = generation_model.generate(
        inputs["input_ids"],
        max_length=150,
        temperature=0.7,
        top_k=50,
        top_p=0.9,
        num_return_sequences=1
    )

    # 解码生成的序列
    generated_answer = tokenizer_generation.decode(output_sequences[0], skip_special_tokens=True)
    return generated_answer


# 完整RAG系统流程
def RAG_system(query):
    # Step 1: 检索模块获取相关内容
    contexts = retrieve_context(query)
    print("检索到的上下文信息:")
    for i, context in enumerate(contexts, 1):
        print(f"{i}. {context}")

    # Step 2: 生成模块生成答案
    answer = generate_answer(contexts, query)
    return answer


# 测试查询
query_text = "人工智能的应用有哪些？"
answer = RAG_system(query_text)

print("\n最终生成的回答:")
print(answer)
